package belf.migrate.training;

import com.belf.framework.web.domain.server.Sys;
import com.google.gson.Gson;
import com.google.gson.JsonArray;

import com.google.gson.JsonObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 从大模型中查询promt，返回可用的SQL语句
 */
@Slf4j
public class LLMQueryTool {
    // 大模型调用地址。在控制台可以灵活设置该值
    private static String endpoint = "http://192.168.0.60/v1/chat/completions";
    // 创建 OkHttpClient 对象
    private static OkHttpClient client = new OkHttpClient.Builder()
            .connectTimeout(60, TimeUnit.SECONDS) // 连接超时时间，这里设置为 10 秒
            .readTimeout(300, TimeUnit.SECONDS) // 读取超时时间，这里设置为 15 秒
            .writeTimeout(300, TimeUnit.SECONDS) // 写入超时时间，这里设置为 15 秒
            .build();

    /**
     * 在参数model指定的大模型中查询提示词prompt，返回可执行的SQL语句
     * @param model 大模型名称
     * @param prompt 提示词
     * @return 一系列可执行的SQL语句
     */
    public static List<String> query(String model, String prompt) {
        long startTime = System.currentTimeMillis();
        List<String> results = null;
        // 使用 Gson 创建 JSON 对象
        JsonObject rootObject = new JsonObject();
        rootObject.addProperty("model", model);

        JsonArray messages = new JsonArray();
        JsonObject systemNode = new JsonObject();
        systemNode.addProperty("role", "system");
        systemNode.addProperty("content", "你是一名经验丰富的数据库DBA.");
        messages.add(systemNode);

        JsonObject useNode = new JsonObject();
        useNode.addProperty("role", "user");
        useNode.addProperty("content", prompt);
        messages.add(useNode);
        rootObject.add("messages", messages);
        // llm的temperature参数越低，输出越接近真实情况。在控制台可以灵活设置该值
        rootObject.addProperty("temperature", 0.1f);
        // llm的top_p参数越低，输出越接近真实情况。一般temperature和top_p参数二选一即可。在控制台可以灵活设置该值
        rootObject.addProperty("top_p", 0.3f);
        // 输出 JSON 字符串
        Gson gson = new Gson();
        String json = gson.toJson(rootObject);
        // 设置请求体为 JSON 格式
        RequestBody requestBody = RequestBody.create(json, MediaType.parse("application/json"));

        try {
            // 创建 POST 请求
            Request request = new Request.Builder()
                    .url(endpoint)
                    .post(requestBody)
                    .build();
            // 发送请求并获取响应
            Response response = client.newCall(request).execute();
            // 打印响应内容
            if (response.isSuccessful()) {
                results = parseChatgptResponse(response.body().string(), model);
            }
            response.close();
        } catch (Exception e) {
            log.error("调用大模型生成SQL失败：{}", e.getMessage());
            e.printStackTrace();
        }

        long endTime = System.currentTimeMillis();
        String seconds = String.format("%.2f seconds", (endTime-startTime)/1000.0);
        log.info("Query LLM {} spend {}秒", model, seconds);
        return results;

    }
    private static List<String> parseChatgptResponse(String jsonString, String model) {
        List<String> results = new ArrayList<>();
        // 解析大模型输出 JSON 字符串
        Gson gson = new Gson();
        JsonObject rootObject = gson.fromJson(jsonString, JsonObject.class);
        JsonArray choices = rootObject.getAsJsonArray("choices");

        if (null != choices && choices.size() > 0) {
            JsonObject jsonObject = choices.get(0).getAsJsonObject();
            JsonObject message = jsonObject.get("message").getAsJsonObject();
            final String content = message.get("content").getAsString();
            // 从大模型接口返回的content中提取所有的SQL语句
            int indexStart = 0;
            int indexEnd = 0;
            boolean flag = true;
            String START_TAG = "```sql";
            String END_TAG = "```";
            while(flag) {
                indexStart = content.indexOf(START_TAG, indexStart);
                if (indexStart >= 0) {
                    indexStart += START_TAG.length();
                    indexEnd = content.indexOf(END_TAG, indexStart);
                    if (indexEnd > 0) {
                        String sql = content.substring(indexStart, indexEnd);
                        results.add(sql);
                        //检查索引是否已经越界
                        indexStart = indexStart + END_TAG.length();
                        if (indexStart >= content.length()) {
                            break;
                        }
                    } else {
                        flag = false;
                    }
                } else {
                    flag = false;
                }
            }
        } //所有的SQL都提取完成

        return results;
    }

}
